/*
 Xern Entity Relocation Network
 Copyright (C) 2011 Ronie Salgado <roniesalg@gmail.com>

 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
 in the Software without restriction, including without limitation the rights
 to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 copies of the Software, and to permit persons to whom the Software is
 furnished to do so, subject to the following conditions:

 The above copyright notice and this permission notice shall be included in
 all copies or substantial portions of the Software.

 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
*/

#include <boost/bind.hpp>
#include <boost/date_time/gregorian/gregorian.hpp>
#include "Xern/DataStream.hpp"
#include "Xern/Log.hpp"
#include "GridNode.hpp"
#include "IMessage.hpp"
#include "IMessageDecoder.hpp"
#include "ControlMessage.hpp"

namespace XernProtocol
{
	const posix_time::time_duration GridNode::NetworkThreadInterval = posix_time::milliseconds(20);
	const posix_time::time_duration GridNode::PacketResendInterval = posix_time::milliseconds(200);
	const posix_time::time_duration GridNode::KeepAliveInterval = posix_time::milliseconds(2000);
	const posix_time::time_duration GridNode::SessionTimeout = posix_time::milliseconds(3000);
	const posix_time::time_duration GridNode::PacketDuplicationTolerance = posix_time::milliseconds(2000);
	const posix_time::ptime GridNode::PacketEpoch(boost::gregorian::date(2000, boost::gregorian::Jan, 1), posix_time::seconds(0));
	const posix_time::time_duration GridNode::ConnectTimeout = posix_time::milliseconds(200);
	const int GridNode::ConnectTries = 3;
	const int GridNode::MaxResendCount = 4;

	GridNode::GridNode()
		: nodeSocket(ioService)
	{
		isInitialized = false;
		isShuttingDown = false;
		packetSequence = 0;
		sessionSequence = 0;
	}

	GridNode::~GridNode()
	{
		Shutdown();
	}

	void GridNode::Initialize()
	{
		boost::unique_lock<boost::mutex> l(nodeMutex);
		if(isInitialized)
			return;

		// Reset the shutting down flag.
		isShuttingDown = false;

		// Create the send thread.
		boost::thread st(boost::bind(&GridNode::SendThread, this));
		sendThread.swap(st);

		// Create the receive thread.
		boost::thread rt(boost::bind(&GridNode::ReceiveThread, this));
		receiveThread.swap(rt);

		// Set initialized flag.
		isInitialized = true;
	}

	void GridNode::Shutdown()
	{
		boost::unique_lock<boost::mutex> l(nodeMutex);
		if(!isInitialized)
			return;

		// Tell the threads to stop.
		isShuttingDown = true;
		ioService.stop();
		l.unlock();

		// Wait the threads to end.
		sendThread.join();
		receiveThread.join();

		isInitialized = false;
	}

	u32 GridNode::GetPacketSequence()
	{
		boost::unique_lock<boost::mutex> l(packetSequenceMutex);
		return ++packetSequence;
	}

	u32 GridNode::GetSessionSequence()
	{
		boost::unique_lock<boost::mutex> l(sessionSequenceMutex);
		return ++sessionSequence;
	}

	void GridNode::SendThread()
	{
		// Lock the node.
		nodeMutex.lock();

		while(!isShuttingDown)
		{
			// Free cpu time.
			// TODO: Use an adaptive system.
			nodeMutex.unlock();
			boost::this_thread::sleep(NetworkThreadInterval);
			nodeMutex.lock();

			// Iterate through each one of the sessions.
			LocalSessions::iterator it = localSessions.begin();
			while(it != localSessions.end())
			{
				// Give the session some time.
				boost::shared_ptr<Session> session = it->second;
				Session_Think(session);

				// Remove disconnected sessions.
				if(session->state == Session::SS_DISCONNECTED)
					localSessions.erase(it++);
				else
					it++;
			}
		}

		nodeMutex.unlock();
	}

	void GridNode::ReceiveThread()
	{
		// Start receiving packets
		nodeSocket.async_receive_from(boost::asio::buffer(receivedPacket.packetData), receiveEndpoint,
				boost::bind(&GridNode::OnPacketReceived, this,
						boost::asio::placeholders::error,
						boost::asio::placeholders::bytes_transferred));

		// Pass the control to the I/O service.
		ioService.run();
	}

	void GridNode::OnPacketReceived(const boost::system::error_code &err, size_t size)
	{
		if(!err)
		{
			// Process the packet
			receivedPacket.packetSize = size;
			ProcessReceivedPacket(receivedPacket, receiveEndpoint);

			// Continue receiving.
			nodeSocket.async_receive_from(boost::asio::buffer(receivedPacket.packetData), receiveEndpoint,
					boost::bind(&GridNode::OnPacketReceived, this,
							boost::asio::placeholders::error,
							boost::asio::placeholders::bytes_transferred));
		}
	}

	void GridNode::ProcessReceivedPacket(Packet &packet, udp::endpoint &endpoint)
	{
		boost::unique_lock<boost::mutex> l(nodeMutex);

		// Decode the packet.
		packet.DecodePacket();

		// Find a matching session.
		LocalSessions::const_iterator it = localSessions.find(packet.destinationSession);
		if(it != localSessions.end())
		{
			boost::shared_ptr<Session> session = it->second;

			// Lock the session.
			boost::unique_lock<boost::mutex> sl(session->sessionMutex);

			// Connecting sessions haven't a valid remote session id.
			if(session->remoteEndpoint == endpoint &&
			   session->state == Session::SS_CONNECTING)
			{
				if(packet.flags & Packet::FL_ACCEPT)
				{
					// Session connection opened.
					session->state = Session::SS_CONNECTED;
					session->remoteSessionId = packet.sourceSession;
					connectionCondition.notify_all();

					// Store the remote session data.
					RemoteSessionId remoteId(endpoint, session->remoteSessionId);
					remoteSessions.insert(std::make_pair(remoteId, session));
				}
				else if(packet.flags & Packet::FL_DENIED)
				{
					// Session connection refused.
					session->state = Session::SS_DISCONNECTED;
					connectionCondition.notify_all();
					return;
				}
			}

			// Check if the packet is valid.
			if(session->remoteSessionId == packet.sourceSession &&
				session->remoteEndpoint == endpoint)
			{
				// Process the packet.
				Session_ReceivePacket(session, packet);
				return;
			}
		}

		// Check for session opening.
		if(packet.flags & Packet::FL_OPEN)
		{
			// Open a new session.
			Session_Open(packet, endpoint);
		}
	}

	void GridNode::Session_Opened(boost::shared_ptr<Session> session)
	{
	}

	void GridNode::Session_Timeout(boost::shared_ptr<Session> session)
	{
	}

	void GridNode::Session_Closed(boost::shared_ptr<Session> session)
	{
	}

	void GridNode::Session_ReceivePacket(boost::shared_ptr<Session> session, Packet &packet)
	{
		// Update the last receive time.
		session->lastReceiveTime = posix_time::microsec_clock::universal_time();

		// Check for session end.
		if(packet.flags & Packet::FL_CLOSE)
		{
			Session_Close(session);
			return;
		}

		// Check for reliable packets.
		if(packet.flags & Packet::FL_RELIABLE)
		{
			// Acknowledge later the packet.
			session->packetsToAck.insert(packet.sequenceNumber);
		}

		// Check for duplicated packets.
		Session::ReceivedPackets::iterator it =
				session->receivedPackets.find(packet.sequenceNumber);
		if(it != session->receivedPackets.end())
		{
			// Update the receive time.
			it->second = posix_time::microsec_clock::universal_time();
			return;
		}

		// Check for appended acknowledges.
		if(packet.flags & Packet::FL_ACKS)
		{
			// Read the appended acknowledges.
			DataStream in(packet.ackCount*4, packet.ackData);
			for(size_t i = 0; i < packet.ackCount; i++)
			{
				u32 packetId;
				in >> packetId;

				// Erase the packet from the waiting list.
				session->packetsWaitingAck.erase(packetId);
			}
		}

		// Packets without payload are control only.
		if(packet.payloadSize == 0)
			return;

		// Decode the message and store in a queue.
		boost::shared_ptr<IMessage> message = session->messageDecoder->DecodeMessage(&packet);
		if(message)
			session->receiveQueue.push_back(message);
	}

	void GridNode::Session_Open(Packet &packet, const udp::endpoint &endpoint)
	{
		// Avoid opening the same session multiple times.
		RemoteSessionId remoteId(endpoint, packet.sourceSession);
		RemoteSessions::const_iterator it = remoteSessions.find(remoteId);
		if(it != remoteSessions.end())
			return;

		// Create a new session.
		boost::shared_ptr<Session> session(new Session());

		// Set the local session id.
		session->sessionId = GetSessionSequence();

		// Set the remote session id.
		session->remoteEndpoint = endpoint;
		session->remoteSessionId = packet.sourceSession;

		// Set the session state.
		session->state = Session::SS_CONNECTED;

		// Set the default message decoder.
		session->messageDecoder = defaultMessageDecoder;

		// Store the new session.
		localSessions.insert(std::make_pair(session->sessionId, session));
		remoteSessions.insert(std::make_pair(remoteId, session));

		// Send a connection established message.
		boost::shared_ptr<ControlMessage> message(new ControlMessage());
		message->flags = Packet::FL_RELIABLE | Packet::FL_ACCEPT;
		message->channel = 0xFF; // Control channel.
		session->Send(message);

		// Notify the lower hierarchy.
		Session_Opened(session);
	}

	void GridNode::Session_Close(boost::shared_ptr<Session> session)
	{
		// Update the session state.
		session->state = Session::SS_DISCONNECTED;

		// Remove the session.
		remoteSessions.erase(RemoteSessionId(session->remoteEndpoint, session->remoteSessionId));

		// Notify the lower hierarchy.
		Session_Closed(session);
	}

	void GridNode::Session_Think(boost::shared_ptr<Session> session)
	{
		// Lock the session.
		boost::unique_lock<boost::mutex> l(session->sessionMutex);

		switch(session->state)
		{
		case Session::SS_CONNECTED:
			Session_Send(session);
			Session_Resend(session);
			Session_Acknowledge(session);
			Session_KeepAlive(session);
			Session_Clean(session);
			break;
		case Session::SS_CONNECTING:
			Session_Send(session);
			break;
		default:
		case Session::SS_DISCONNECTED:
			break;
		}
	}

	void GridNode::Session_Send(boost::shared_ptr<Session> session)
	{
		while(!session->sendQueue.empty())
		{
			// Get the now time.
			posix_time::ptime now = posix_time::microsec_clock::universal_time();

			// Remove the front message.
			boost::shared_ptr<IMessage> message = session->sendQueue.front();
			session->sendQueue.pop_front();

			// TODO: Take into account the throttling.

			// Create the packet.
			boost::shared_ptr<Packet> packet(new Packet());

			// Write the packet header.
			packet->flags = message->GetFlags();
			packet->sequenceNumber = GetPacketSequence();
			packet->sourceSession = session->sessionId;
			packet->destinationSession = session->remoteSessionId;
			packet->channel = message->GetChannel();
			packet->sentTime = (now - PacketEpoch).total_milliseconds();
			packet->lastSentTime = now;
			packet->WriteHeader();

			// Write the packet payload.
			message->WriteToPacket(packet.get());

			// TODO: Try to append acks instead of creating a dedicated packet.

			// Update the last send time.
			session->lastSendTime = now;

			// Send the packet.
			nodeSocket.send_to(boost::asio::buffer(packet->packetData, packet->packetSize),
					session->remoteEndpoint);

			// Store reliable packets.
			if(packet->flags & Packet::FL_RELIABLE)
			{
				session->packetsWaitingAck.insert(std::make_pair(packet->sequenceNumber, packet));
			}

			// Check for close packets.
			if(packet->flags & Packet::FL_CLOSE)
			{
				// Set the new session state.
				session->state = Session::SS_DISCONNECTED;

				// Remove the session.
				remoteSessions.erase(RemoteSessionId(session->remoteEndpoint, session->remoteSessionId));

				// Notify the other threads.
				disconnectionCondition.notify_all();

				// Don't continue sending packets.
				return;
			}
		}
	}

	void GridNode::Session_Resend(boost::shared_ptr<Session> session)
	{
		// Calculate the now time.
		posix_time::ptime now = posix_time::microsec_clock::universal_time();

		// Iterate through each packet
		Session::PacketsWaitingAck::iterator it = session->packetsWaitingAck.begin();
		while(it != session->packetsWaitingAck.end())
		{
			boost::shared_ptr<Packet> packet = it->second;
			if(now < packet->lastSentTime + PacketResendInterval)
			{
				it++;
				continue;
			}

			// Check the packet resend count.
			if(packet->resendCount == MaxResendCount)
			{
				// Remove the packet.
				session->packetsWaitingAck.erase(it++);
				continue;
			}

			// Increase the packet resend count.
			packet->flags |= Packet::FL_RESENT;
			packet->resendCount++;
			packet->lastSentTime = now;
			//packet->UpdateHeader();

			// Update the last send time.
			session->lastSendTime = now;

			// Resend the packet.
			nodeSocket.send_to(boost::asio::buffer(packet->packetData, packet->packetSize),
								session->remoteEndpoint);

			// Increase the iterator.
			it++;
		}
	}

	void GridNode::Session_Acknowledge(boost::shared_ptr<Session> session)
	{
		// Check if there are acks waiting.
		if(session->packetsToAck.empty())
			return;

		// Copy the packet ids.
		std::vector<u32> packetIds;
		int numacks = session->packetsToAck.size();
		packetIds.reserve(numacks);

		Session::PacketsToAck::iterator it = session->packetsToAck.begin();
		for(; it != session->packetsToAck.end(); it++)
			packetIds.push_back(*it);
		session->packetsToAck.clear();

		// Calculate the number of packets to send.
		int numpackets = numacks / 255;
		if(numacks % 255 != 0)
			numpackets ++;

		// Create the packet.
		Packet packet;
		packet.flags = Packet::FL_ACKS;
		packet.channel = 0xFF;
		packet.sourceSession = session->sessionId;
		packet.destinationSession = session->remoteSessionId;
		packet.lastSentTime = posix_time::microsec_clock::universal_time();
		packet.sentTime = (packet.lastSentTime - PacketEpoch).total_milliseconds();

		// Send the packets.
		int pos = 0;
		for(int i = 0; i < numpackets; i++)
		{
			// Write the packet header.
			packet.sequenceNumber = GetPacketSequence();
			packet.WriteHeader();

			// Write the payload.
			packet.WritePayload(0, NULL);

			// Write the acks.
			int toPut = numacks >= 255 ? 255 : numacks;
			packet.WriteAcks(toPut, &packetIds[pos]);

			// Increase the acks position.
			pos += toPut;

			// Send the packet.
			nodeSocket.send_to(boost::asio::buffer(packet.packetData, packet.packetSize),
					session->remoteEndpoint);
		}
	}

	void GridNode::Session_KeepAlive(boost::shared_ptr<Session> session)
	{
		// Check if its time to send a keep alive.
		posix_time::ptime now = posix_time::microsec_clock::universal_time();
		if(now < session->lastSendTime + KeepAliveInterval)
			return;

		// Send the keep-alive.
		boost::shared_ptr<ControlMessage> message(new ControlMessage());
		message->flags = Packet::FL_RELIABLE;
		message->channel = 0xFF;
		session->sendQueue.push_back(message);
	}

	void GridNode::Session_Clean(boost::shared_ptr<Session> session)
	{
		// Clean messages in the received packets list.
		posix_time::ptime now = posix_time::microsec_clock::universal_time();
		Session::ReceivedPackets::iterator it = session->receivedPackets.begin();
		while(it != session->receivedPackets.end())
		{
			if(now < it->second + PacketDuplicationTolerance)
			{
				it++;
				continue;
			}

			// Erase the old packet.
			session->receivedPackets.erase(it++);
		}

		// Check for session timeout.
		if(session->lastReceiveTime + SessionTimeout < now)
		{
			// Notify the lower hierarchy.
			Session_Timeout(session);

			// Remove the session.
			session->state = Session::SS_DISCONNECTED;
			remoteSessions.erase(RemoteSessionId(session->remoteEndpoint, session->remoteSessionId));
		}

	}

	boost::shared_ptr<Session> GridNode::ConnectSession(const udp::endpoint &endpoint)
	{
		// Lock the node.
		boost::unique_lock<boost::mutex> lock(nodeMutex);

		// Create a new session.
		boost::shared_ptr<Session> session(new Session);
		session->messageDecoder = defaultMessageDecoder;
		session->sessionId = GetSessionSequence();
		session->remoteEndpoint = endpoint;
		session->state = Session::SS_CONNECTING;

		// Register the session.
		localSessions.insert(std::make_pair(session->sessionId, session));

		// Try to connect a limited number of times.
		for(int i = 0; i < ConnectTries; i++)
		{
			// Create the connection message.
			boost::shared_ptr<ControlMessage> message(new ControlMessage());
			message->flags = Packet::FL_OPEN;
			message->channel = 0xFF;

			// Send the message.
			session->Send(message);

			// Wait until a response.
			connectionCondition.timed_wait(lock, ConnectTimeout);

			// Check the new session state.
			if(session->state == Session::SS_CONNECTED)
				break;
		}

		// Remove the timeout session.
		if(session->state == Session::SS_CONNECTING)
		{
			localSessions.erase(session->sessionId);
			LogConsole->Error("Session timeout.");
		}
		else if(session->state == Session::SS_DISCONNECTED)
		{
			LogConsole->Error("Session refused.");
		}

		return session->state == Session::SS_CONNECTED ? session :
			boost::shared_ptr<Session> ();
	}

	void GridNode::DisconnectSession(boost::shared_ptr<Session> session, bool wait)
	{
		if(!session)
			return;

		// Send a session close message.
		boost::unique_lock<boost::mutex> lock(nodeMutex);

		// Create the close message
		boost::shared_ptr<ControlMessage> message(new ControlMessage());
		message->flags = Packet::FL_CLOSE;
		message->channel = 0xFF;

		// Close the session.
		session->Send(message);

		// Wait until the message is sent.
		if(wait)
			disconnectionCondition.timed_wait(lock, ConnectTimeout);
	}
}; // namespace XernProtocol
